#include <cassert>
#include <iostream>
#include <fstream>
#include <string>
#include <cmath>
#include <ctime>
#include <algorithm>

#include "sampler.h"
#include "str_util.h"
#include "model.h"

// n_status: number of latent status for each biterm
void Model::init() {
  cout << "init" << endl;
  
  nwz.resize(K1, W);
  phi_z.resize(K1, W); // uniformly init
  nb_z.resize(K1);

  nwc.resize(K0, W);
  phi_c.resize(K0, W); // uniformly init
  nb_c.resize(K0);

  nt0 = 0;
  nt1 = 0;
  // random initialize
  for (int i = 0; i < bs.size(); ++i) {
	double pt_b = bs[i].get_pt_b();
	int t = Sampler::Bern_sample(pt_b);
	int k;
	if (t == 0)
	  k = Sampler::uni_sample(K0);
	else
	  k = Sampler::uni_sample(K1);

	assign_biterm_topic(bs[i], t, k);
  }
}

// input, each line is a doc
// format: wid  wid  burstiness
void Model::load_biterms() {
  cout << "load biterms: " << dfile << endl;
  ifstream rf( dfile.c_str() );
  if (!rf) {
	cout << "file not find:" << dfile << endl;
	exit(-1);
  }

  string line;
  while(getline(rf, line)) {
	Biterm bi(line);
	bs.push_back(bi);
  }
  cout << "all bs num:" << bs.size() << endl;
}

void Model::run() {
  load_biterms();
  init();
  
  cout << "Begin iteration" << endl;
  for (int it = 1; it < n_iter + 1; ++it) {
	cout << "iter " << it << endl;

	for (int b = 0; b < bs.size(); ++b) {
	  update_biterm(bs[b]);
	}
	
	cout << " nt1/N = " << nt1 << "/" << bs.size() 
		 << " = " << double(nt1) / bs.size() << endl;

	if (it % save_step == 0)
	  save_all();
  }

  save_all();
}

// sample procedure for ith biterm 
double Model::update_biterm(Biterm& bi) {
  reset_biterm_topic(bi);
  // compute p(z|b)
  Pvec<double> pz;
  compute_pz_b(bi, pz);
  // sample topic for biterm b
  int k = Sampler::mult_sample(pz.to_vector());
  int t = 1;
  if (k < K0) 
	t = 0;
  else 
	k = k - K0;
  assign_biterm_topic(bi, t, k);
  return pz.max();
}

// reset topic assignment of biterm i
void Model::reset_biterm_topic(Biterm& bi) {
  int w1 = bi.get_wi();
  int w2 = bi.get_wj();
  int t = bi.get_t();
  int k = bi.get_z();

  if (t == 0) {
	nb_c[k] -= 1;	// update number of biterms in topic K
	nwc[k][w1] -= 1;	// update w1's occurrence times in topic K
	nwc[k][w2] -= 1;
	nt0 -= 1;
	assert(nb_c[k] >= 0 && nwc[k][w1] >= 0 && nwc[k][w2] >= 0);
  }
  else {
	nb_z[k] -= 1;	// update number of biterms in topic K
	nwz[k][w1] -= 1;	// update w1's occurrence times in topic K
	nwz[k][w2] -= 1;
	nt1 -= 1;
	assert(nb_z[k] >= 0 && nwz[k][w1] >= 0 && nwz[k][w2] >= 0);
  }
  
  bi.reset_z();
}

// compute p(z|w_i, w_j)
void Model::compute_pz_b(Biterm& bi, Pvec<double>& pz) {
  pz.resize(K0 + K1);
  int w1 = bi.get_wi();
  int w2 = bi.get_wj();
  double pt_b = bi.get_pt_b();

  double pw1k, pw2k, pk_t;
  // common topic
  for (int k = 0; k < K0; ++k) {
	double deno = W / (2 * nb_c[k] + W * beta0);
	pw1k = (nwc[k][w1] + beta0) * deno;
	pw2k = (nwc[k][w2] + beta0) * deno;
	pk_t = (nb_c[k] + alpha0) / (nt0 + K0 * alpha0);
	pz[k] = (1-pt_b) * pk_t * pw1k * pw2k;
  }

  // bursty topic
  for (int k = 0; k < K1; ++k) {
	double deno = W / (2 * nb_z[k] + W * beta1);
	pw1k = (nwz[k][w1] + beta1) * deno;
	pw2k = (nwz[k][w2] + beta1) * deno;
	pk_t = (nb_z[k] + alpha1) / (nt1 + K1*alpha1);
	pz[k + K0] = pt_b * pk_t * pw1k * pw2k;
  }

  pz.normalize();
}

// assign topic k to biterm i
void Model::assign_biterm_topic(Biterm& bi, int t, int k) {
  int w1 = bi.get_wi();
  int w2 = bi.get_wj();

  bi.set_z(t, k);
  if (t == 0) {
	nb_c[k] += 1;
	nwc[k][w1] += 1;
	nwc[k][w2] += 1;
	nt0 += 1;
  }
  else {
	nb_z[k] += 1;
	nwz[k][w1] += 1;
	nwz[k][w2] += 1;
	nt1 += 1;
  }
}

void Model::save_all() {
  save_pz();
  save_pc();
  save_pw_z();
  save_pw_c();
  save_btz();
}

// p(z) is determinated by the overall proportions
// of biterms in it
void Model::save_pz() {
  Pvec<double> pz(K1);	          // p(z) = theta
  for (int k = 0; k < K1; k++) 
	pz[k] = (nb_z[k] + alpha1);
  pz.normalize();
  
  string pt = dir + "pz";
  cout << "writing p(z): " << pt << endl;
  pz.write(pt);
}

void Model::save_pc() {
  Pvec<double> pc(K0);	          // p(z) = theta
  for (int k = 0; k < K0; k++) 
	pc[k] = (nb_c[k] + alpha0);
  pc.normalize();
  
  string pt = dir + "pc";
  cout << "writing p(c): " << pt << endl;
  pc.write(pt);
}

void Model::save_pw_z() {
  Pmat<double> pw_z(K1, W, beta1);   // p(w|z) = phi, size K * M
  pw_z += nwz;
  pw_z.normr();

  string pt = dir + "pw_z";
  cout << "writing p(w|z): " << pt << endl;
  //pw_z.write(pt);
  nwz.write(pt);
}

void Model::save_pw_c() {
  Pmat<double> pw_c(K0, W, beta0);   // p(w|z) = phi, size K * M
  pw_c += nwc;
  pw_c.normr();

  string pt = dir + "pw_c";
  cout << "writing p(w|c): " << pt << endl;
  pw_c.write(pt);
}

void Model::save_btz() {
  string pt = dir + "btz";
  cout << "writing btz: " << pt << endl;
  ofstream wf(pt.c_str());
  for (int i = 0; i < bs.size(); ++i) {
	wf << bs[i].str() << endl;
  }
}
